I don't think this is the correct logic. Asymptotically for an infinitely deep context it should always be worthwhile to run as many CUDA blocks as possible since the overhead becomes negligible. Intuitively I would expect something like this to be a better solution: if possible, always run at least 2 blocks / SM in order to keep the GPU busy when calling __syncthreads(). After that, only increase the number of blocks / SM while each block gets some minimum chunk of the KV cache to work on.

I don't think this is the correct logic. Asymptotically for an infinitely deep context it should always be worthwhile to run as many CUDA blocks as possible since the overhead becomes negligible. Intuitively I would expect something like this to be a better solution: if possible, always run at least 2 blocks / SM in order to keep the GPU busy when calling __syncthreads(). After that, only increase the number of blocks / SM while each block gets some minimum chunk of the KV cache to work on.

I agree with what you are saying. But in practice, we are seeing good speed-up even for models that have max_blocks_per_sm = 2 and this PR reduces it to 1 block per SM. It's possible if we continue increasing context length, the perf benefits might reduce. I will try to collect that data.

I will also explore if there is a better way to reduce the overhead of flash_attn_stream_k_fixup without reducing blocks per SM.

An RTX 5090 and an RTX Pro 6000 both have 192 SMs with Qwen 3 30b a3b with 4 KV heads that will be on average a 171 chunk / SM. The kernel will run with an internal batch size (nbatch_fa) of either 64 or 128 depending on Q->ne[1]. It may make sense to try increasing the minimum chunk size / SM to something like 256.

As for the stream-k fixup itself: if the occupancy is low anyways we could maybe run the kernel with a number of CUDA blocks that is an exact multiple of ntiles_dst and write an alternative fixup kernel with simplified logic that will maybe be faster.

An RTX 5090 and an RTX Pro 6000 both have 192 SMs with Qwen 3 30b a3b with 4 KV heads that will be on average a 171 chunk / SM. The kernel will run with an internal batch size (nbatch_fa) of either 64 or 128 depending on Q->ne[1]. It may make sense to try increasing the minimum chunk size / SM to something like 256.

I'm not sure if this will help in all cases. As you can see in the perf data, this change is helping even for 32K context length, where ntiles_KV*ntiles_dst is much bigger than max_blocks. In such cases, the chunk size per SM should already be large enough.

As for the stream-k fixup itself: if the occupancy is low anyways we could maybe run the kernel with a number of CUDA blocks that is an exact multiple of ntiles_dst and write an alternative fixup kernel with simplified logic that will maybe be faster.

Thanks for the idea. I will look into it.

Closing this PR in favor of #21159